#spatial weight

import numpy as np

# rainfall: 2D array, shape = (m, n)

# m = number of processes, n = number of grids
# rainfall[i, j] = rainfall of the j-th grid in the i-th process
a = np.sum(rainfall)  # total rainfall across all grids and processes (scalar)
cij = rainfall * (rainfall / a)  # shape = (m, n)
c_j = np.sum(cij, axis=0)  # sum across processes, shape = (n,)
# c_j[j] = weighted total contribution of the j-th grid across all processes
c_total = np.sum(c_j)  # scalar, should equal np.sum(cij)
spatial_weight = c_j / c_total  # shape = (n,), sum = 1
